import cv2
import io
import numpy as np
import matplotlib.pyplot as plt


def imgaussflpf(image, sigma):
    rows, cols = image.shape[:2]
    x, y = np.meshgrid(
        np.linspace(-cols / 2, cols / 2, cols), np.linspace(-rows / 2, rows / 2, rows)
    )
    gaussian_filter = np.exp(-((x**2 + y**2) / (2 * sigma**2)))
    return gaussian_filter


def imfreqfilter(image_in, ff):
    if len(image_in.shape) == 3 and image_in.shape[2] == 3:
        image_in = np.dot(image_in[..., :3], [0.2989, 0.5870, 0.1140])
    if image_in.shape != ff.shape:
        raise ValueError("Image and filter are not the same size.")

    f = np.fft.fft2(image_in)
    s = np.fft.fftshift(f)
    image_out = np.fft.ifftshift(s * ff)
    image_out = np.abs(np.fft.ifft2(image_out))
    return image_out / np.max(image_out)


def gauss_filter(image_stream):
    nparr = np.frombuffer(image_stream.read(), np.uint8)
    I1 = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    # 如果为彩色图像转换为灰色图像
    if len(I1.shape) == 3:
        I1 = cv2.cvtColor(I1, cv2.COLOR_BGR2GRAY)

    # 生成滤镜
    ff1 = imgaussflpf(I1, 20)
    ff2 = imgaussflpf(I1, 40)
    ff3 = imgaussflpf(I1, 60)

    # 应用滤镜
    B1 = imfreqfilter(I1, ff1)
    B2 = imfreqfilter(I1, ff2)
    B3 = imfreqfilter(I1, ff3)

    fig, axs = plt.subplots(2, 2, figsize=(10, 8))

    axs[0, 0].imshow(I1, cmap="gray")
    axs[0, 0].set_title("Original Image")

    axs[0, 1].imshow(B1, cmap="gray")
    axs[0, 1].set_title("20 Gaussian Low-pass Filtered Image")

    axs[1, 0].imshow(B2, cmap="gray")
    axs[1, 0].set_title("40 Gaussian Low-pass Filtered Image")

    axs[1, 1].imshow(B3, cmap="gray")
    axs[1, 1].set_title("60 Gaussian Low-pass Filtered Image")
    plt.tight_layout()
    plt.show()

    # 返回图像的字节数据（这里假设你是想以某种方式返回图像数据，实际应用中可以根据需求调整）
    # 需要注意的是，这里的方式可能不是最适合你的需求的，你可以根据实际情况调整
    buffer = io.BytesIO()
    plt.savefig(buffer, format="png")
    buffer.seek(0)
    return buffer.read()


if __name__ == "__main__":
    gauss_filter("../test1_3.jpg")
